import os
import PIL
from torchvision import datasets, transforms
import torchvision
from timm.data import create_transform

# CIFAR10的均值和标准差
CIFAR10_MEAN = (0.4914, 0.4822, 0.4465)
CIFAR10_STD = (0.2023, 0.1994, 0.2010)

def build_dataset(is_train, args):
    transform = build_transform(is_train, args)
    
    # 直接使用CIFAR10类，不需要指定train/val子目录
    dataset = datasets.CIFAR10(
        root=args.data_path,  # 直接使用根目录
        train=is_train,       # 通过这个参数区分训练集和测试集
        download=True,        # 如果数据集不存在就下载
        transform=transform
    )
    
    print(f"Dataset CIFAR10 {'train' if is_train else 'test'} size: {len(dataset)}")
    return dataset

def build_transform(is_train, args):
    # 使用CIFAR10的均值和标准差
    mean = CIFAR10_MEAN
    std = CIFAR10_STD
    
    if is_train:
        # 训练集变换
        transform = transforms.Compose([
            transforms.RandomCrop(32, padding=4),  # CIFAR10标准做法
            transforms.RandomHorizontalFlip(),     # 随机水平翻转
            transforms.Resize(args.input_size),    # 调整到目标大小
            transforms.ToTensor(),
            transforms.Normalize(mean, std),
            transforms.RandomErasing(p=args.reprob) if args.reprob > 0 else None,
        ])
        # 移除None项
        transform.transforms = [t for t in transform.transforms if t is not None]
        return transform
    
    else:
        # 测试集变换
        return transforms.Compose([
            transforms.Resize(args.input_size),  # 调整到目标大小
            transforms.ToTensor(),
            transforms.Normalize(mean, std)
        ])
